# author: lgx
# date: 2022-10-19 12:48:15
# description:
from torch.utils.data import Dataset
from torchvision.transforms import *
import copy
from pipeline import Compose


class MyDataset(Dataset):
    def __init__(self, transform_dict=None, **kwargs):
        super(MyDataset, self).__init__(**kwargs)
        self.transform_dict = transform_dict
        self.transform = self.init_transform()
        self.data_infos = self._get_data_info()

    def _get_data_infos(self):
        return []

    def __len__(self):
        return len(self.data_infos)

    def init_transform(self):
        if self.transform_dict is None:
            return None
        transform_dict = self.transform_dict.get(self.phase, None)
        if transform_dict is None:
            return None
        transform = Compose(transform_dict)
        return transform

    def __getitem__(self, item, debug=False):
        """Prepare the frames for training given the index."""
        results = copy.deepcopy(self.data_infos[item])
        """
            some code to add keywords
        """
        return self.transform(results) if self.transform is not None else results

    def _get_input(self, item):
        info = self.data_infos[item]
        input = 0
        if self.transform is not None:
            input = self.transform(input)
        return input

    def _generate_target(self, item):
        info = self.data_infos[item]
        target = 0
        return target

    # used for debug
    def visualize(self, item, save_path=None):
        input, target = self.__getitem__(item)
        # visualize code

